import os
import numpy as np
import pandas as pd
from PIL import Image
from Models import Multi_models


# 排序{“汉字”：“序列=1,2,3...”}
def label_of_directory(directory):
    """
    sorted for label indices
    return a dict for {'classes', 'range(len(classes))'}
    """
    classes = []
    # 遍历每一个分类文件夹
    for subdir in sorted(os.listdir(directory)):
        # 判断分类文件夹下是否有图片
        if os.path.isdir(os.path.join(directory, subdir)):
            # 将分类加入列表
            classes.append(subdir)

    class_indices = dict(zip(classes, range(len(classes))))
    return class_indices


# 从字典的value获得key --> 如：字典{“白”：1}
def get_key_from_value(dict, index):
    for keys, values in dict.items():
        if values == index:
            return keys


# 产生图片路径列表
def generator_list_of_imagepath(path):
    image_list = []
    for image in os.listdir(path):
        if not image == '.DS_Store':
            image_list.append(path + image)
    return image_list


# 读取数据
def load_image(image):
    img = Image.open(image)
    img = img.resize((128, 128))
    img = np.array(img)
    img = img / 255
    # reshape img to size(1, 128, 128, 1)
    img = img.reshape((1,) + img.shape + (1,))
    return img


# 预测图片，得到概率最高的那个分类
def get_label_predict_top1(image, model):
    """
    image = load_image(image), input image is a ndarray
    retturn best of label
    """
    predict_proprely = model.predict(image)
    predict_label = np.argmax(predict_proprely, axis=1)
    return predict_label


# 预测图片，得到概率最高的top_k个分类
def get_label_predict_top_k(image, model, top_k):
    """
    image = load_image(image), input image is a ndarray
    return top-5 of label
    """
    # array 2 list
    predict_proprely = model.predict(image)
    predict_list = list(predict_proprely[0])
    min_label = min(predict_list)
    label_k = []
    for i in range(top_k):
        label = np.argmax(predict_list)
        predict_list.remove(predict_list[label])
        predict_list.insert(label, min_label)
        label_k.append(label)
    return label_k


# 得出test数据集中的top1 --> precision
def test_image_predict_top1(model, test_image_path, directory):

    model.load_weights(WEIGHTS_PATH) # 导入训练好的权重参数
    image_list = generator_list_of_imagepath(test_image_path) # 测试图片路径导入

    predict_label = []
    class_indecs = label_of_directory(directory) # 得到分类的字典{class：num}
    for image in image_list:
        img = load_image(image)
        label_index = get_label_predict_top1(img, model) # 预测的key
        label = get_key_from_value(class_indecs, label_index) # 从字典中得到key，也就是汉字
        predict_label.append(label)

    return predict_label


# 得出test数据集中的top5 --> recall
def test_image_predict_top_k(model, test_image_path, directory, top_k):

    model.load_weights(WEIGHTS_PATH)
    image_list = generator_list_of_imagepath(test_image_path)

    predict_label = []
    class_indecs = label_of_directory(directory)
    for image in image_list:
        img = load_image(image)
        # return a list of label max->min
        label_index = get_label_predict_top_k(img, model, top_k)
        label_value_dict = []
        for label in label_index:
            label_value = get_key_from_value(class_indecs, label)
            label_value_dict.append(str(label_value))

        predict_label.append(label_value_dict)

    return predict_label


# 将列表转化为str类型的标签
def tran_list2str(predict_list_label):
    new_label = []
    for row in range(len(predict_list_label)):
        str = ""
        for label in predict_list_label[row]:
            str += label
        new_label.append(str)
    return new_label


# 保存为csv文件
def save_csv(test_image_path, predict_label):
    image_list = generator_list_of_imagepath(test_image_path)
    save_arr = np.empty((10000, 2), dtype=np.str)
    save_arr = pd.DataFrame(save_arr, columns=['filename', 'lable'])
    predict_label = tran_list2str(predict_label)
    for i in range(len(image_list)):
        filename = image_list[i].split('/')[-1]
        save_arr.values[i, 0] = filename
        save_arr.values[i, 1] = predict_label[i]
    save_arr.to_csv('submit_test.csv', decimal=',', encoding='utf-8', index=False, index_label=False)
    print('submit_test.csv have been write, locate is :', os.getcwd())


if __name__ == "__main__":
    # print(dict(zip(["a", "b"], range(2))))
    # 训练好的模型参数路径
    WEIGHTS_PATH = "logs/best_mobile.h5"
    # 训练数据的路径
    train_path = './data/train_data/'
    # test数据路径
    test_image_path = './data/test1/test1/'

    print("=====test label=====")
    mmodels = Multi_models(CLASS=100, size=128)
    simple_model = mmodels.mobile()
    simple_model.load_weights(WEIGHTS_PATH)
    model = simple_model
    predict_label = test_image_predict_top_k(model, test_image_path, train_path, 5)
    print(predict_label)

    print("======csv save======")
    save_csv(test_image_path, predict_label)

    print("finished！！！")